# install required packages
!pip install dashvector dashscope
!pip install transformers_stream_generator python-dotenv
# prepare news corpus as knowledge source
!git clone https://github.com/shijiebei2009/CEC-Corpus.git
import dashscope
import os
from dotenv import load_dotenv
from dashscope import TextEmbedding
from dashvector import Client, Doc

# get env variable from .env
# please make sure DASHSCOPE_KEY is defined in .env
load_dotenv()
dashscope.api_key = os.getenv('DASHSCOPE_KEY')


# initialize DashVector for embedding's indexing and searching
dashvector_client = Client(api_key='{your-dashvector-api-key}')

# define collection name
collection_name = 'news_embeddings'

# delete if already exist
dashvector_client.delete(collection_name)

# create a collection with embedding size of 1536
rsp = dashvector_client.create(collection_name, 1536)
collection = dashvector_client.get(collection_name)
def prepare_data_from_dir(path, size):
    # prepare the data from a file folder in order to upsert to DashVector with a reasonable doc's size.
    batch_docs = []
    for file in os.listdir(path):
        with open(path + '/' + file, 'r', encoding='utf-8') as f:
            batch_docs.append(f.read())
            if len(batch_docs) == size:
                yield batch_docs[:]
                batch_docs.clear()

    if batch_docs:
        yield batch_docs
def prepare_data_from_file(path, size):
    # prepare the data from file in order to upsert to DashVector with a reasonable doc's size.
    batch_docs = []
    chunk_size = 12
    with open(path, 'r', encoding='utf-8') as f:
        doc = ''
        count = 0
        for line in f:
            if count < chunk_size and line.strip() != '':
                doc += line
                count += 1
            if count == chunk_size:
                batch_docs.append(doc)
                if len(batch_docs) == size:
                    yield batch_docs[:]
                    batch_docs.clear()
                doc = ''
                count = 0

    if batch_docs:
        yield batch_docs
def generate_embeddings(docs):
    # create embeddings via DashScope's TextEmbedding model API
    rsp = TextEmbedding.call(model=TextEmbedding.Models.text_embedding_v1,
                             input=docs)
    embeddings = [record['embedding'] for record in rsp.output['embeddings']]
    return embeddings if isinstance(docs, list) else embeddings[0]
id = 0
dir_name = 'CEC-Corpus/raw corpus/allSourceText'

# indexing the raw docs with index to DashVector
collection = dashvector_client.get(collection_name)

# embedding api max batch size
batch_size = 4  

for news in list(prepare_data_from_dir(dir_name, batch_size)):
    ids = [id + i for i, _ in enumerate(news)]
    id += len(news)
    # generate embedding from raw docs
    vectors = generate_embeddings(news)
    # upsert and index
    ret = collection.upsert(
        [
            Doc(id=str(id), vector=vector, fields={"raw": doc})
            for id, doc, vector in zip(ids, news, vectors)
        ]
    )
    print(ret)
# check the collection status
collection = dashvector_client.get(collection_name)
rsp = collection.stats()
print(rsp)
def search_relevant_context(question, topk=1, client=dashvector_client):
    # query and recall the relevant information
    collection = client.get(collection_name)

    # recall the top k similarity results from DashVector
    rsp = collection.query(generate_embeddings(question), output_fields=['raw'],
                           topk=topk)
    return "".join([item.fields['raw'] for item in rsp.output])
# query the top 1 results
question = '清华博士发生了什么？'
context = search_relevant_context(question, topk=1)
print(context)
2006-08-26 10:41:45
8月23日上午9时40分，京沪高速公路沧州服务区附近，一辆由北向南行驶的金杯面包车撞到高速公路护栏上，车上5名清华大学博士后研究人员及1名司机受伤，被紧急送往沧州二医院抢救。截至发稿时，仍有一名张姓博士后研究人员尚未脱离危险。
# initialize qwen 7B model
from modelscope import AutoModelForCausalLM, AutoTokenizer
from modelscope import GenerationConfig

tokenizer = AutoTokenizer.from_pretrained("qwen/Qwen-7B-Chat", revision = 'v1.0.5',trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("qwen/Qwen-7B-Chat", revision = 'v1.0.5',device_map="auto", trust_remote_code=True, fp16=True).eval()
model.generation_config = GenerationConfig.from_pretrained("Qwen/Qwen-7B-Chat",revision = 'v1.0.5', trust_remote_code=True) # 可指定不同的生成长度、top_p等相关超参
# define a prompt template for the vectorDB-enhanced LLM generation
def answer_question(question, context):
    prompt = f'''请基于```内的内容回答问题。"
	```
	{context}
	```
	我的问题是：{question}。
    '''
    history = None
    print(prompt)
    response, history = model.chat(tokenizer, prompt, history=None)
    return response
# test the case on plain LLM without vectorDB enhancement
question = '清华博士发生了什么？'
answer = answer_question(question, '')
print(f'question: {question}\n' f'answer: {answer}')
请基于```内的内容回答问题。"
	```
	
	```
	我的问题是：清华博士发生了什么？。
    
question: 清华博士发生了什么？
answer: 清华博士是指清华大学的博士研究生。作为一名AI语言模型，我无法获取个人的身份信息或具体事件，因此无法回答清华博士发生了什么。如果您需要了解更多相关信息，建议您查询相关媒体或官方网站。
# test the case with knowledge
context = search_relevant_context(question, topk=1)
answer = answer_question(question, context)
print(f'question: {question}\n' f'answer: {answer}')
请基于```内的内容回答问题。"
	```
	2006-08-26 10:41:45
8月23日上午9时40分，京沪高速公路沧州服务区附近，一辆由北向南行驶的金杯面包车撞到高速公路护栏上，车上5名清华大学博士后研究人员及1名司机受伤，被紧急送往沧州二医院抢救。截至发稿时，仍有一名张姓博士后研究人员尚未脱离危险。


	```
	我的问题是：清华博士发生了什么？。
    
question: 清华博士发生了什么？
answer: 8月23日上午9时40分，一辆由北向南行驶的金杯面包车撞到高速公路护栏上，车上5名清华大学博士后研究人员及1名司机受伤，被紧急送往沧州二医院抢救。